import numpy as np
import cv2
import os
class genStrokes(object):
     def __init__(self, config, ):
		 
          self.num = config.batch_sz * config.M
          self.im_sz = config.img_sz
          self.im_ch = config.img_channels
          
          self.window_size = np.array((config.win_size, config.win_size))
          self.threathhold = 0.99
          self.step_width = 5
          self.step_height = 5
          self.strokeColorBg = (0, 0, 0)       # black stroke
          self.strokeColorFg = (255, 255, 255) # white stroke00
     
     def genStrokesForOne(self, alphaImg, rgbImg):
          img_width = alphaImg.shape[0]
          img_height = alphaImg.shape[1]
          num_pt_bg = 0
          num_pt_fg = 0
          for i in range(img_width):
              for j in range(img_height):
                  if alphaImg[i,j] == 0:
                        num_pt_bg += 1
                  elif alphaImg[i,j] == 255:
                        num_pt_fg += 1
          num_pt_all = img_height * img_width
          Bgratio = num_pt_bg / num_pt_all
          Fgratio = num_pt_fg / num_pt_all
          if Bgratio > self.threathhold:
              rgbImg[:, :] = self.strokeColorBg
          elif Fgratio > self.threathhold:
              rgbImg[:, :] = self.strokeColorFg
          return rgbImg.copy()
     
     

     def genStroke(self, strokeMap, loc, rgbImg, alphaImg):
          leftup = [max(0,(loc[0] - self.window_size[0]/2)),max(0,(loc[1] - self.window_size[1]/2))]
          
          rightdown = [min(self.im_sz,(loc[0] + self.window_size[0]/2)),min(self.im_sz,(loc[1] + self.window_size[1]/2))]
          
          temp_alpha_img = alphaImg[leftup[1]:rightdown[1], leftup[0]:rightdown[0]]
          temp_rgb_img = rgbImg[leftup[1]:rightdown[1],leftup[0]:rightdown[0]]
          
          strokes_img = self.genStrokesForOne(temp_alpha_img,temp_rgb_img)
          strokeMap[leftup[1]:rightdown[1], leftup[0]:rightdown[0]] = strokes_img
          return strokeMap


     def drawstroke(self, strokemap, loc, gtAlpha, counter):
          strokeMap = strokemap
          loc = loc           # [1, 2]   N = batch*M
          image = strokemap
          gtAlpha = gtAlpha
          cnt = counter
          
          
          loc_this = loc
          img_shape = np.array((self.im_sz, self.im_sz))
          loc_this = (loc_this * img_shape[-1::-1]/2.0 + img_shape[-1::-1]/2.0).astype(np.int32)
          image_this = image
          gtAlpha_this = gtAlpha
          strokeMap_this = strokemap
          strokeMap_tmp = self.genStroke(strokeMap_this, loc_this, image_this, gtAlpha_this)
          #strokeMap_tmp = np.expand_dims(strokeMap_tmp, 0)
          return strokeMap_tmp
               

          
          
